Skip to content

refactor(core): migrate DatagenStore onto StorageBase - #217

Merged
beinan merged 2 commits into
lance-format:mainfrom
beinan:datagen-store-base
Jul 28, 2026
Merged

refactor(core): migrate DatagenStore onto StorageBase#217
beinan merged 2 commits into
lance-format:mainfrom
beinan:datagen-store-base

Conversation

@beinan

@beinan beinan commented Jul 27, 2026

Copy link
Copy Markdown
Collaborator

Third and last store of #214 Part 1, after #215 and #216. All three fixed-schema stores now share one storage implementation.

Stacked on #216 — review that first; this diff is datagen_store.rs only.

Bugs this closes

Compaction, which this store had none of. Every WAL merge appends a fragment to the base table, and datagen seals once per checkpoint batch, so fragments grew without bound and scans degraded with them. There was no compact_files call anywhere in the file. Now compact / should_compact / compaction_stats, with defer_index_remap set — the fieldless MemWAL index panics Lance's inline remap.

Scalar index, which it also had none of. DatasetIndexExt was imported solely for load_indices in mem_wal_index_present. create_event_id_index now builds a ZoneMap on event_id; point lookups by event id previously always scanned.

load_with_options dropped storage options. It fell back to a bare Dataset::open whenever no options were passed, and never attached a session — so every flushed-generation read allocated Lance's 6 GiB default cache, keyed per generation URI.

Deleted duplication

merge_own_shard was a near-line-for-line copy of rollout's, and buffered every flushed generation fully in memory. Gone, along with merge_own_shard_if_ready, the writer lifecycle (write_with_resident_writer / put_seal_drain / ensure_write_writer), ensure_mem_wal, mem_wal_index_present, shard-snapshot discovery, flushed_generation_uri, open_flushed_dataset, load_with_options, create_with_options, the Drop impl, and a second copy of is_fenced_error.

Net -269 lines in this file.

Behavior preserved exactly

seal_on_put: true. Each append is one complete checkpoint batch that must be durable and readable before the writer moves on, so a crash cannot expose a partially checkpointed step. This is the store's existing putforce_seal_activewait_for_flush_drain behavior, unchanged.

key_column is event_id, not id. Event ids are derived deterministically (UUIDv5 over item/checkpoint/ordinal), which is what makes a retried append idempotent under LSM dedup. The base already took the key column as a parameter, so it needed no change to accommodate this.

Verification

New test compaction_folds_merged_fragments_and_preserves_reads pins behavior that could not exist before: fragments accumulate across merges, compaction reduces the count, and every event stays readable afterward.

  • 181 core lib tests pass (up from 180)
  • 182 Python tests pass
  • cargo clippy --workspace --all-targets clean

What's left in #214

Part 1 is done after this — all three stores share StorageBase. Part 2 (arbitrary user-defined schemas with a required id) builds on it: the base already validates key_column against the schema at open, which is the enforcement point that work needs.

🤖 Generated with Claude Code

beinan and others added 2 commits July 27, 2026 22:37
Second step of lance-format#214, after lance-format#215. `ContextStore` now delegates its storage
layer to `StorageBase` instead of carrying its own forked implementation.

Closes the gaps that made this store diverge from rollout:

- **WAL merge, which this store never had.** Flushed generations accumulated
  under `_mem_wal/` forever and every read unioned all of them, so read cost
  grew without bound in the number of writes. It now gets the same merge as
  rollout, exposed as `maybe_merge_wal` (count trigger),
  `cleanup_wal` (time trigger) and `pending_wal_generations`.
- **Compaction now defers index remap.** The base table carries a fieldless
  MemWAL index and Lance's inline remap panics on it.
- **Compaction and index creation no longer drop `storage_options`.** Both
  reloaded via a bare `Dataset::open`, which silently discarded credentials
  and broke on any authenticated object store. Both now reload through
  `StorageBase::load_with_options`.
- **Resident writer.** `add` used to open, `put` and `close` a `ShardWriter`
  on every call; it now reuses one resident writer with fence-retry, so an
  append no longer pays a cold DNS resolution + TCP/TLS handshake + epoch
  claim.
- **`add` takes `&self`.** The writer lives behind the base's mutex, so the
  server takes a read lock and concurrent appends stop serializing on the
  store's `RwLock`.

Sharding changes to match rollout: writes go to the shard owned by this
writer instance (`ContextStoreOptions::shard_id`) rather than one derived
per `(bot_id, session_id)`. One instance owns exactly one shard, so the
epoch-fencing invariant holds by construction, and merge has a single
well-defined shard to drain. `derive_region_id` is removed.

`add` visibility is preserved, not silently weakened. `StorageBase` gains a
`seal_on_put` switch: rollout keeps its deferred seal (durable on return,
visible via the flush sweeper or `?flush=true`), while `ContextStore`
defaults to `seal_on_add: true` — required for correctness, since its
id/external-id uniqueness validation, upsert and tombstones all read the
table back. Callers appending trusted batches can opt into rollout's
throughput profile with `seal_on_add: false`.

`ContextStore` is no longer `Clone`: it owns a `StorageBase`, which owns the
resident writer and a `Drop` that seals it, so it must have exactly one
owner. The two clone sites now open their own handle instead — the
background compactor (which only rewrites base-table fragments and would
otherwise contend with the write path) and Python's `Context.fork` (which
branches the in-memory context over the same dataset).

Three new tests pin the behavior: `add` is visible on return by default,
deferred seal is durable-but-invisible until `flush`, and WAL generations
merge into the base table and stay readable. Full workspace suite passes
(180 core lib tests, up from 177), clippy clean.

Co-Authored-By: Claude <noreply@anthropic.com>
Third and last store of lance-format#214 Part 1, after lance-format#215 and lance-format#216. All three
fixed-schema stores now share one storage implementation.

Closes this store's gaps:

- **Compaction, which it had none of.** Every WAL merge appends a fragment
  to the base table and datagen seals once per checkpoint batch, so
  fragments grew without bound and scans degraded with them. Now exposed as
  `compact` / `should_compact` / `compaction_stats`, with `defer_index_remap`
  set (the fieldless MemWAL index panics Lance's inline remap).
- **Scalar index, which it also had none of.** `create_event_id_index`
  builds a ZoneMap on `event_id`; point lookups by event id previously
  always scanned.
- **`load_with_options` dropped storage options.** It fell back to a bare
  `Dataset::open` whenever no options were passed, and never attached a
  session, so every generation read allocated Lance's 6 GiB default cache.

Deletes the near-line-for-line copy of rollout's `merge_own_shard` (plus
`merge_own_shard_if_ready`, the writer lifecycle, `ensure_mem_wal`, shard
snapshot discovery, the `Drop` impl and a second `is_fenced_error`), and
picks up the streaming-safe version in the base instead of buffering every
flushed generation in memory.

`seal_on_put: true`: each append is one complete checkpoint batch that must
be durable *and* readable before the writer moves on, so a crash cannot
expose a partially checkpointed step. This preserves the store's existing
put -> seal -> drain behavior exactly.

`key_column` is `event_id`, not `id` — event ids are derived
deterministically (UUIDv5 over item/checkpoint/ordinal), which is what makes
a retried append idempotent under LSM dedup. The base already took the key
column as a parameter, so no change was needed there.

New test `compaction_folds_merged_fragments_and_preserves_reads` pins the
behavior that did not exist before: fragments accumulate across merges,
compaction reduces them, and every event stays readable.

Full suite passes — 181 core lib tests, 182 Python tests, clippy clean.

Co-Authored-By: Claude <noreply@anthropic.com>
@beinan
beinan merged commit 1a1b6b2 into lance-format:main Jul 28, 2026
9 checks passed
beinan added a commit that referenced this pull request Jul 28, 2026
Part 2 of #214, on top of the consolidated `StorageBase` from #215#217.

Users can now declare their own schema instead of picking one of the
three built-in ones — while `add` and the WAL merge behave identically
to `RolloutStore`, because they are literally the same code path, not a
reimplementation.

## Three pieces

### `schema_spec` — schema as a value

The built-in stores hard-code an Arrow schema and then repeat its field
list across five to seven places (schema fn, builders, decoder,
projection, filters, merge key, index name). A `SchemaSpec` is declared
once at creation and persisted with the store.

Validation enforces what the storage layer actually assumes:

- **`id` is mandatory** — `Utf8`, non-nullable, tagged
`unenforced-primary-key`. It is always the LSM merge key and the indexed
column, which is what makes a retried append idempotent and a crashed
merge safe to redo.
- **Reserved names rejected** — anything starting with `_`, so future
Lance internals can't collide with a user column.
- **Vector columns carry dim + metric**, so an index can be built and
the metric round-trips on open without being re-specified.
- **Nested lists and lists-of-vectors rejected** rather than silently
mis-encoded.

### `generic_codec` — schema-driven encode/decode

JSON rows ↔ Arrow, matched **by name in both directions**. Two
deliberate choices:

- An undeclared key is an **error, not a silent drop** — a field-name
typo fails at the write instead of becoming missing data discovered much
later.
- Projected batches decode without their missing columns, which is what
lets a blob-excluding scan round-trip.

This also fixes by construction a bug class the built-in stores have:
they decode nested structs (`relationships`, `state_metadata`)
**positionally**, so reordering fields silently mis-assigns data.

### `generic_store` — the store

Delegates *all* storage behavior to `StorageBase`: resident writer,
fence retry, surgical WAL drain, compaction with `defer_index_remap`, id
index, LSM reads. `seal_on_add` picks rollout's deferred-seal profile
(default) or read-your-write.

## How users write it

```python
store = GenericStore.open(uri, schema={
    "id":        {"type": "string", "nullable": False},
    "user_id":   {"type": "string"},
    "score":     {"type": "float32"},
    "embedding": {"type": "vector", "dim": 768, "metric": "cosine"},
    "video":     {"type": "binary", "blob": True},
})

store.add([{"id": "r1", "user_id": "u1", "score": 0.9}])   # nullable cols may be omitted
store.list()                                  # no video — never materialized
store.get("r1", columns=["video"])            # fetch the blob for one row
```

## Blobs are inline, and that's deliberate

Per your point that this project exists to store big blobs: **inline
`LargeBinary`, never blob-v2 offload**. The LSM read path has no
blob-materialization step, so an offloaded column reads back as `None` —
which is exactly why `RolloutStore` stores artifacts inline today.

`blob: true` doesn't change *where* bytes live, it changes *who pays for
them*: blob columns are dropped from default scan projections, so `list`
never materializes them, and `get(id, columns)` fetches them per row.

I measured this path directly before designing around it:

| size | write | read | survives WAL merge |
|---|---|---|---|
| 5 MB | 80 ms | 42 ms | ✓ |
| 40 MB | 483 ms | 262 ms | ✓ |
| 120 MB | 1.39 s | 760 ms | ✓ |

The test suite pins a 4 MB round-trip through a merge (kept small so CI
stays fast).

⚠️ **Worth knowing for the 100 MB+ case:** `max_memtable_size` defaults
to 256 MB, so two such rows trigger a seal. That's not breakage — the
count/time merge triggers bound the resulting generations — but it does
raise seal frequency. `ShardWriterConfig` exposes the knob; no store
surfaces it yet. Follow-up if it bites.

## Schema is immutable in v1

Reopening with a different schema is **rejected**, not silently
reinterpreted over existing data. Evolution needs a versioning story and
is out of scope here.

## Verification

28 new tests: 8 spec validation, 9 codec round-trip, 11 store behavior —
including that duplicate ids dedup to the newest write (proving `id`
really is the merge key), deferred seal matches rollout's visibility
semantics, and generations merge into the base table.

209 core lib tests pass (up from 181), `cargo clippy --workspace
--all-targets` clean.

## Not in this PR

The server/client/Python surface — `CreateContextRequest` carrying a
schema, generic record DTOs, deduping the six duplicated conversion
functions, explicit column names in search/retrieve. This PR is the core
engine; the API layer is the next one, and it's large enough to want its
own review.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

---------

Co-authored-by: Claude <noreply@anthropic.com>
beinan added a commit that referenced this pull request Jul 28, 2026
The three follow-ups from #220, plus a regression guard for a bug I
introduced while writing them.

## 1. `seal_on_add` is now persisted in the dataset

It lived only in `GenericStoreOptions`, so a store created with
`seal_on_add: true` **lost read-your-write the first time it was evicted
from the LRU and reopened** — the server passed a hardcoded `false` at
`state.rs:642`.

The failure mode is the nasty kind: no error, no data loss, just
"sometimes I can't read what I just wrote" — and triggered by *cache
pressure*, not by anything in the caller's code. Untestable locally,
only shows up under load. `/merge-wal` would accidentally paper over it
(it flushes internally), making it look intermittent.

The mode now sits in the schema metadata beside the spec, so a reopen
restores what the store was *created* with. Stores written before this
fall back to the caller's option, so nothing breaks on upgrade.

## 2. The sweepers now visit every store kind

Both `spawn_flush_sweeper` and `spawn_global_sweeper` hardcoded
`state.rollout_stores`. Consequences:

- **datagen**: seals on every append, so flush didn't matter — but
*nothing ever merged its generations*. They accumulated forever and
every read unioned all of them. This predates #217.
- **generic**: defaults to deferred seal, so writes stayed invisible
until someone called `/flush` by hand.

Rather than copy the loop a third time — the exact duplication #214
removed one layer down — the traversal is generic over a new `Sweepable`
trait in `sweeper.rs`.

**Why the trait is on `Arc<RwLock<Store>>` and not on the store:**
rollout's merge deliberately splits into a shared-lock prepare (the
expensive object-storage reads, during which appends keep flowing) and a
brief exclusive-lock commit. A trait over `&mut Store` would have forced
the exclusive lock across the whole merge and quietly stalled the write
path — a performance regression that no test would have caught. Letting
each kind own its locking preserves that.

Metric names keep their `rollout_` prefix and gain a `kind` label, so
existing dashboards and alerts keep working.

## 3. Config naming: documented, not renamed

`ROLLOUT_FLUSH_INTERVAL_SECS` and friends now govern all three kinds, so
the prefix is misleading. But renaming breaks deployment manifests for
zero functional gain, so I documented the actual scope in the flag docs
and the startup warning instead. Happy to add neutral aliases if you'd
rather.

## 4. ⚠️ A regression I caused, and the guard for it

While rewriting those doc comments my edit deleted two `#[arg(..)]`
attributes, turning `--rollout-flush-interval-secs` and
`--rollout-cleanup-interval-secs` into **required positional
arguments**:

```
error: the following required arguments were not provided:
  <ROLLOUT_CLEANUP_INTERVAL_SECS>
  <ROLLOUT_FLUSH_INTERVAL_SECS>
```

Everything still compiled. All 203 Rust tests still passed. The server
just refused to start. It surfaced only because the Python remote tests
spawn the real binary — and those tests swallow stderr, so it showed up
as "server did not become healthy".

Two tests in `config.rs` now close that hole: every field must be an
optional flag with a default, and the minimal documented invocation must
parse.

## Verification

I checked each fix actually fails without its fix, rather than trusting
a green run:

- reverting the persistence → both `seal_mode_survives_*` tests fail
- dropping the `#[arg]` → both config tests fail

9 new tests, including e2e ones that go through the real server path:
seal mode surviving an actual LRU eviction
(`generic_stores.lock().pop()`), and a sweeper pass draining a real
datagen generation.

Full suite: **203 core + 62 server + 18 api** Rust, **199 Python**.
clippy, ruff, pyright all clean.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-authored-by: Claude <noreply@anthropic.com>
beinan pushed a commit that referenced this pull request Jul 28, 2026
… layers (#219)

## What

Completes the `DatagenStore` surface the datagen POC needs, wired
through all six layers (core → api → client/server/unified → PyO3 →
python wrapper) and exposed on the high-level
`lance_context.DatagenStore` wrapper so callers using the public Python
API reach every new function.

## Pieces

- **A1 — `DatagenStreamWriter`**: a client-side state machine that
stamps the bookkeeping columns (`item_seq`, `attempt`, `checkpoint_id`,
`event_id`) and returns event dicts to hand to
`append`/`append_checkpoint`. No new HTTP endpoint — works embedded and
remote. Fresh streams via `open_stream` (attempt=0, next_seq=1); resume
via `resume_stream` (attempt=`last_attempt+1`, continuing `item_seq`).
- **A3 — resume / attempt>0 fold semantics**: `resuming_writer` rebuilds
a writer from the folded item; the fold folds correctly across attempts.
- **A4 — `load_blob` by field name**: core `load_blob(folded,
field_name)` resolves the folded item's `blob_event_ids` map to bytes.
Propagated to Python via the folded dict's `blob_event_ids` +
`get_blob`, composed in the `api.py` wrapper.
- **A5 — inspection tree**: `item_tree` folds every projected descendant
to its latest state and links parent→child. Server does a thin raw
`events_for_root` dump; the client folds the tree via a single-source
`DatagenItemTree::build`.

## Python API

The high-level `lance_context.DatagenStore` wrapper gains `open_stream`
/ `resume_stream` / `item_tree` / `load_blob`, plus a
`DatagenStreamWriter` wrapper class (`step_started`, `step_completed`,
`item_terminal`, `item_failed`, `item_id`/`attempt` properties). Both
are re-exported from `lance_context`.

## Tests

- `python/tests/test_datagen.py` — 6 tests covering
open/resume/item_tree/load_blob/item_failed through the public wrapper.
- A core store integration test driving a real store: open a root
stream, complete a step with a Set field, checkpoint, terminal, spawn a
child stream, then assert `item_tree` links
roots/children/status/query_tags/fields.

All green: Rust 25 datagen core tests, Python 6 tests, `ruff`,
`pyright`, `cargo fmt --check`.

## Note

Branches from before origin/main merged the StorageBase migration
(#216/#217/#218). CI may surface conflicts against those; will
rebase/fix as needed.

---------

Co-authored-by: YangjunZ <yangjunzhang@microsoft.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant